package ai.mspbots.sql.ai;

import ai.mspbots.sql.constant.Prompt;
import ai.mspbots.sql.dto.QueryDTO;
import cn.hutool.core.util.StrUtil;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.HttpEntity;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpMethod;
import org.springframework.http.HttpStatus;
import org.springframework.http.ResponseEntity;
import org.springframework.stereotype.Service;
import org.springframework.web.client.RestClientException;
import org.springframework.web.client.RestTemplate;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

@Service
@Slf4j
public class DeepSeekOperation {

    private static final String API_KEY = "sk-9eec7d1b37a748828bc58568f05d5ba8";

    // DeepSeek API 的端点
    private static final String DEEPSEEK_ENDPOINT = "https://api.deepseek.com/chat/completions";
    @Autowired
    private RestTemplate restTemplate;

    @Autowired
    private ObjectMapper objectMapper;

    public List<QueryDTO> explainSql(String sql) {
        String text = getJsonNodeResponseEntity(StrUtil.format(Prompt.EXPLAIN_PROMPT, sql));
        List<QueryDTO> queryDTOList = new ArrayList<>();
        try {
            JsonNode jsonNode = objectMapper.readTree(text);
            for (int i = 0; i < jsonNode.size(); i++) {
                JsonNode node = jsonNode.get(i);
                queryDTOList.add(new QueryDTO(node.get("table").asText(), node.get("query").asText(), node.get("count").asText()));
            }
        } catch (JsonProcessingException e) {
            log.error("Failed to parse JSON response from DeepSeek API", e);
            throw new RuntimeException(e);
        }
        return queryDTOList;
    }


    public String formatSql(String sql) {
        return getJsonNodeResponseEntity(StrUtil.format(Prompt.FORMAT_SQL, sql));
    }

    private String getJsonNodeResponseEntity(String prompt) {
        HttpHeaders headers = new HttpHeaders();
        headers.set("Content-Type", "application/json");
        headers.set("Authorization", "Bearer " + API_KEY);

        Map<String, Object> requestMap = new HashMap<>();
        requestMap.put("model", "deepseek-chat");
        requestMap.put("stream", false);

        // 创建 messages 列表
        List<Map<String, String>> messages = new ArrayList<>();

        // 添加 user 消息
        Map<String, String> userMessage = new HashMap<>();
        userMessage.put("role", "user");
        userMessage.put("content", prompt);
        messages.add(userMessage);

        // 将 messages 添加到 requestMap
        requestMap.put("messages", messages);


        HttpEntity<Object> entity = new HttpEntity<>(requestMap, headers);

        try {
            ResponseEntity<JsonNode> response = restTemplate.exchange(DEEPSEEK_ENDPOINT, HttpMethod.POST, entity, JsonNode.class);
            if (response.getStatusCode() == HttpStatus.OK) {
                JsonNode body = response.getBody();
                // 从 DeepSeek API 的响应中提取 assistant 的回复
                String text = body.at("/choices/0/message/content").asText();
                return text.replaceAll("```json", "").replaceAll("```", "");
            }
        } catch (RestClientException e) {
            throw new RuntimeException(e);
        }
        throw new RuntimeException("Failed to get response from DeepSeek API");
    }

}
